package com.ideal.gpt.ability.docqa.splitter;

import java.util.ArrayList;
import java.util.List;

import com.lzwork.utility.file.LineContent;
import com.lzwork.utility.file.LineContent.LoadSetting;
import com.lzwork.utility.file.LineContent.LoadSetting.SkipLineType;
import com.lzwork.utility.pair.StringPair;
import com.lzwork.utility.pair.StringTuple;
import com.lzwork.utility.str.StrUtility;

/**
 * @author LiZheng
 * @date 2023年9月11日 下午4:57:02
 */

public class PairPrefix
{
	public PairPrefix(String q, String a, boolean firstOnly, boolean allowNoA)
	{
		headQ = q;
		headA = a;
		this.firstOnly = firstOnly;
		this.allowNoA = allowNoA;
	}

	String headQ;
	String headA;
	int mode = 0;
	boolean firstOnly;
	boolean allowNoA;

	@Override
	public String toString()
	{
		return "[" + headQ + " - " + headA + "] first only:" + firstOnly + ", allow no answer:" + allowNoA;
	}

	public int mode()
	{
		return mode;
	}

	public int checkMode(String firstLine)
	{
		if (firstLine.startsWith(headQ))
			mode = 1;
		else if (firstLine.startsWith("1." + headQ))
			mode = 2;
		else if (firstLine.startsWith("1. " + headQ))
			mode = 3;
		else if (firstLine.startsWith("1. "))
			mode = 4;

		return mode;
	}

	StringTuple getHead(int num)
	{
		String head = "";
		if (mode != 1)
			head += num + ".";

		if (mode == 3)
			head += " ";

		if (mode == 4)
			return new StringTuple(head, headA, head);

		return new StringTuple(head + headQ, head + headA, ":");
	}

	String getContext(String line, StringTuple headSp)
	{
		int split = -1;

		if (!StrUtility.stringValid(line))
			return "";

		if (headSp != null)
			split = line.indexOf(headSp.third());

		if (split < 0)
		{
			split = line.indexOf(':');
			if (split >= 0 && split < 10)
				split += 1;
		} else
			split += headSp.third().length();

		if (split >= 0)
		{
			return line.substring(split).trim();
		}

		return line;
	}

	public List<StringPair> splitContentPairs(String input)
	{
		if (!StrUtility.stringValid(input))
			return new ArrayList<>();

		input = input.replaceAll("：", ":");
		String[] lines = input.split("[\n]");
		List<StringPair> result = new ArrayList<>();

		StringBuilder currentQ = new StringBuilder();
		StringBuilder currentA = new StringBuilder();

		int state = 0;

		if (firstOnly)
			checkMode(lines[0]);
		else
		{
			int i = 0;
			while (i < lines.length)
			{
				String line = lines[i];
				i++;
				int m = checkMode(line);
				if (m != 0)
					break;
			}
		}

		if (mode != 0)
		{
			int index = 1;
			StringTuple headSp = getHead(index);

			for (String l : lines)
			{
				String line = l.trim();
				boolean isQ = false;
				boolean isA = false;

				if (state == 0 || state == 2)
				{
					isQ = line.startsWith(headSp.first());
				} else if (state == 1)
				{
					isA = line.startsWith(headSp.second()) || line.startsWith(headA);

					if (!isA && allowNoA)
					{
						StringTuple headSpNext = getHead(index + 1);
						isQ = line.startsWith(headSpNext.first());
						if (isQ)
						{
							headSp = headSpNext;
							index++;
						}
					}
				}

				if (isQ)
				{
					state = 1;

					if (StrUtility.stringValid(currentQ))
					{
						if (StrUtility.stringValid(currentA) || allowNoA)
							result.add(new StringPair(currentQ.toString(), currentA.toString()));
						currentQ = new StringBuilder();
						currentA = new StringBuilder();
					}

					String context = getContext(line, headSp);
					currentQ.append(context);

				} else if (isA)
				{
					state = 2;

					String context = getContext(line, null);
					currentA.append(context);

					index++;
					headSp = getHead(index);
				} else if (state == 1)
				{
					currentQ.append("\n");
					currentQ.append(line);
				} else if (state == 2)
				{
					currentA.append("\n");
					currentA.append(line);
				}
			}

			if (StrUtility.stringValid(currentQ))
			{
				if (StrUtility.stringValid(currentA) || allowNoA)
					result.add(new StringPair(currentQ.toString(), currentA.toString()));
				currentQ = new StringBuilder();
				currentA = new StringBuilder();
			}
		}

		return result;
	}

	public List<StringPair> splitByParagraph(String input)
	{
		input = input.replaceAll("：", ":");
		String qPart = headQ;
		String aPart = headA;

		List<Integer> qIndex = StrUtility.findAllKey(qPart, input);
		List<Integer> aIndex = StrUtility.findAllKey(aPart, input);

		if (qIndex.size() != 1 || aIndex.size() != 1)
			return splitContentPairs(input);

		String qStr = input.substring(qIndex.get(0) + qPart.length(), aIndex.get(0));
		String aStr = input.substring(aIndex.get(0) + aPart.length());

		LoadSetting setting = new LoadSetting();
		setting.setTrimSpace(true);
		setting.setLineSkip(SkipLineType.SkipEmpty);

		List<String> qList = LineContent.splitLines(qStr, setting);
		List<String> aList = LineContent.splitLines(aStr, setting);

		int max = Math.min(qList.size(), aList.size());

		List<StringPair> result = new ArrayList<>();

		for (int i = 0; i < max; i++)
		{
			String q = qList.get(i);
			String a = aList.get(i);

			String pfx = Integer.toString(i + 1) + ".";

			if (q.startsWith(pfx))
				q = q.substring(2).trim();
			if (a.startsWith(pfx))
				a = a.substring(2).trim();
			result.add(new StringPair(q, a));
		}

		return result;
	}
}
